#include "MallocCrash.h"
#include <assert.h>
#include <new>
#include <iostream>
#include "../Utils/LogUtil.h"
#include "../Templates/Alignmen.h"
#include "../Templates/TypeCompatibleBytes.h"


struct MallocPool
{
	uint32 mNumUsed;
	uint32 mMaxUsedIndex;
	uint32 mMaxNumUsed;
	uint32 mTotalNumUsed;

	/** Allocation size for this pool. */
	const uint32 mAllocationSize;

	/** Fixed list of allocations for the specified size for this pool. */
	PtrInfo* mAllocations[MallocCrash::MAX_NUM_ALLOCS_IN_POOL];

	/** Maximum number of allocations that can be made for this pool. */
	const uint32 mMaxNumAllocations;

	/** Memory allocated in the pool and memory used by a fixed array. */
	uint32 mAllocatedMemory;

	MallocPool(const PoolDesc& poolDesc, MallocCrash& outer) :
		mNumUsed(0),
		mMaxUsedIndex(0),
		mMaxNumUsed(0),
		mTotalNumUsed(0),
		mAllocationSize(poolDesc.mSize + MallocCrash::PER_ALLOC_OVERHEAD),
		mMaxNumAllocations(poolDesc.mNumAllocs)
	{
		for (uint32 index = 0; index < mMaxNumAllocations; ++index)
		{
			uint8* newPtr = (uint8*)outer.AllocateFromPool(mAllocationSize);
			mAllocations[index] = new(newPtr) PtrInfo(newPtr + MallocCrash::PER_ALLOC_OVERHEAD);
		}

		// Zero the rest.
		for (uint32 index = mMaxNumAllocations; index < MallocCrash::MAX_NUM_ALLOCS_IN_POOL; ++index)
		{
			mAllocations[index] = nullptr;
		}

		mAllocatedMemory = mMaxNumAllocations * mAllocationSize + sizeof(mAllocations);
	}

	uint8* AllocateFromPool(uint32 allocationSize)
	{
		// Find an unused allocation.
		PtrInfo* ptrInfo = nullptr;
		for (uint32 index = 0; index < mMaxNumAllocations; ++index)
		{
			PtrInfo* ptrIt = mAllocations[index];
			if (ptrIt->mSize == 0)
			{
				ptrInfo = ptrIt;
				mMaxUsedIndex = glm::max<uint32>(mMaxUsedIndex, index);
				break;
			}
		}

		assert(ptrInfo != nullptr);

		mNumUsed++;
		mTotalNumUsed++;
		mMaxNumUsed = glm::max<uint32>(mMaxNumUsed, mNumUsed);
		ptrInfo->mSize = allocationSize;

		memset((void*)ptrInfo->mPtr, MallocCrash::MEM_TAG, ptrInfo->mSize);

		LOGI(L"Malloc Size=%u, PooledPtr=0x%016llx", mAllocationSize, (uint64)ptrInfo->mPtr);

		return ptrInfo->mPtr;
	}

	/** Tries to free a pointer. */
	void TryFreeFromPool(uint8* ptr)
	{
		//const uint32 PtrSize = FGenericPlatformMallocCrash::GetAllocationSize(Ptr);
		//FPlatformMisc::LowLevelOutputDebugStringf( TEXT( "Free SizeWithOverhead=%u, PooledPtr=0x%016llx \n" ), PtrSize, (uint64)Ptr );

		bool isRemoved = false;
		for (uint32 index = 0; index < mMaxNumAllocations; ++index)
		{
			PtrInfo* ptrIt = mAllocations[index];
			if (ptrIt->mPtr == ptr)
			{
				memset((void*)ptrIt->mPtr, MallocCrash::MEM_WIPETAG, ptrIt->mSize);

				ptrIt->mSize = 0;
				isRemoved = true;
				mNumUsed--;

				break;
			}
		}

		assert(isRemoved);
	}

	__forceinline static MallocPool& GetPool(int32_t index)
	{
		assert(index >= 0 && index < MallocCrash::NUM_POOLS);
		return ((MallocPool*)&sBytePools)[index];
	}

private:

	static TTypeCompatibleBytes<MallocPool[MallocCrash::NUM_POOLS]> sBytePools;
};

TTypeCompatibleBytes<MallocPool[MallocCrash::NUM_POOLS]> MallocPool::sBytePools;

MallocCrash::MallocCrash() :
	mMemoryPoolOffset(0)
{
	mMemoryPool = (uint8*)malloc(GetPoolTotalSize());
	assert(mMemoryPool != nullptr);
	for (uint32 index = 0; index < NUM_POOLS; ++index)
	{
		MallocPool* newPool = new(&MallocPool::GetPool(index)) MallocPool(GetPoolDesc(index), *this);
	}
}

MallocCrash::~MallocCrash()
{
	for (int32_t index = NUM_POOLS; index; )
	{
		--index;
		(&MallocPool::GetPool(index))->MallocPool::~MallocPool();
	}
	free(mMemoryPool);
}

void* MallocCrash::Malloc(uint32 count, uint32 alignment /*= DEFAULT_ALIGNMENT*/)
{
	if (alignment > 16)
	{
		assert(0);
	}

	MallocPool* pool = FindPoolFromSize(count);
	if (pool)
	{
		const uint8* pooledPtr = pool->AllocateFromPool(count);
		return (void*)pooledPtr;
	}

	return nullptr;
}

void* MallocCrash::Realloc(void* original, uint32 count, uint32 alignment /*= DEFAULT_ALIGNMENT*/)
{
	assert(original != nullptr);
	if (IsPtrInPool(original))
	{
		uint32 ptrSize = GetAllocationSize(original);
		void* result = Malloc(count, REQUIRED_ALIGNMENT);
		memcpy(result, original, min(count, ptrSize));
		Free(original);
		return result;
	}
	else
	{
		assert(0);
	}

	return nullptr;
}

void MallocCrash::Free(void* original)
{
	if (IsPtrInPool(original))
	{
		MallocPool* pool = FindPoolFromSize(GetAllocationSize(original));
		if (pool)
		{
			pool->TryFreeFromPool((uint8*)original);
		}
		else
		{
			assert(0);
		}
	}
	else
	{
		assert(0);
	}
}

bool MallocCrash::GetAllocationSize(void* original, uint32& sizeOut)
{
	sizeOut = MallocCrash::GetAllocationSize(original);
	return true;
}

bool MallocCrash::IsPtrInPool(void* ptr)
{
	return (ptr >= &mMemoryPool[0] && ptr < &mMemoryPool[GetPoolTotalSize()]);
}

const PoolDesc& MallocCrash::GetPoolDesc(uint32 index) const
{
	static const PoolDesc AllPoolDesc[NUM_POOLS] =
	{
		PoolDesc(64, 224),
		PoolDesc(96, 144),
		PoolDesc(128,  80),
		PoolDesc(192, 560),
		PoolDesc(256, 384),
		PoolDesc(384, 208),
		PoolDesc(512,  48),
		PoolDesc(768,  32),
		PoolDesc(1024,  32),
		PoolDesc(2048,  32),
		PoolDesc(4096,  32),
		PoolDesc(8192,  32),
		PoolDesc(16384,  16),
		PoolDesc(32768,  16)
	};
	return AllPoolDesc[index];
}

uint32 MallocCrash::GetPoolTotalSize() const
{
	static uint32 TotalSize = 0;
	if (TotalSize == 0)
	{
		for (uint32 Index = 0; Index < NUM_POOLS; ++Index)
		{
			const PoolDesc& poolDesc = GetPoolDesc(Index);
			assert(poolDesc.mNumAllocs % 16 == 0);
			assert(poolDesc.mSize % 16 == 0);
			TotalSize += poolDesc.mNumAllocs * (poolDesc.mSize + PER_ALLOC_OVERHEAD);
		}

		// TotalSize = Align(TotalSize, SafePageSize());
		// TotalSize = Align(TotalSize, 0);
	}
	return TotalSize;
}

MallocPool* MallocCrash::FindPoolFromSize(uint32 allocationSize) const
{
	for (uint32 index = 0; index < NUM_POOLS; ++index)
	{
		MallocPool* pool = &MallocPool::GetPool(index);

		// If we have used up this pools allocation limit lets try the next one
		if (pool->mNumUsed >= pool->mMaxNumAllocations)
		{
			continue;
		}

		if (allocationSize <= pool->mAllocationSize - PER_ALLOC_OVERHEAD)
		{
			return pool;
		}
	}

	// Use large allocation pool.
	return nullptr;
}

uint8* MallocCrash::AllocateFromPool(uint32 allocationSize)
{
	if (mMemoryPoolOffset + allocationSize <= GetPoolTotalSize())
	{
		const uint32 returnMemoryPoolOffset = mMemoryPoolOffset;
		mMemoryPoolOffset += allocationSize;
		return (uint8*)mMemoryPool + returnMemoryPoolOffset;
	}
	else
	{
		assert(0);
	}
	return nullptr;
}

uint32 MallocCrash::GetAllocationSize(void* original)
{
	PtrInfo* ptrInfo = (PtrInfo*)((uint8*)original - PER_ALLOC_OVERHEAD);
	return (uint32)(ptrInfo->mSize);
}

